#include <buola/cv/ccloud.h>
#include <buola/scene.h>
#include <buola/scene/geodes/cpointcloud.h>
#include <buola/app/ceventloop.h>
#include <buola/image/io.h>
#include <buola/image/cgc.h>
#include <buola/image/algorithm/filter.h>
#include <buola/app/ccmdline.h>
#include <buola/app/capp.h>
#include <buola/scene/cscene.h>
#include <buola/scene/cimagerenderer.h>
#include <buola/scene/csceneview.h>
#include <buola/scene/cperspectivecamera.h>
#include <buola/scene/transform/crttransform.h>
#include <buola/gui/ctoolkit.h>
#include <buola/scene/geodes/ccube.h>
#include <buola/scene/controllers/ccameracontroller.h>
#include <buola/io/cbinaryreader.h>
#include <buola/algorithm/binary.h>

using namespace buola;

class CROSBag
{
    typedef std::map<std::string,CSimpleBuffer<uint8_t>> THeader;
    
    struct SChunk
    {
        SChunk()=default;
        SChunk(SChunk&&)=default;

        int64_t mOffset;
        int64_t mSize;
        bool mCompressed;
        std::unique_ptr<CSimpleBuffer<uint8_t>> mBuffer;
    };

    struct SMessage
    {
        int mChunk;
        int64_t mTimestamp;
        int32_t mOffset;
    };
    
    struct SConnection
    {
        std::string mTopic;
        std::string mType;
        std::string mDefinition;
        std::vector<SMessage> mMessages;
    };
    
public:
    CROSBag(const io::CURI &pURI)
        :   mReader(pURI)
    {
        ParseBag();
    }
    
    int64_t GetTimestamp(const CSimpleBuffer<uint8_t> &pBuffer)
    {
        int64_t lRaw=binary_interpret<int64_t>(pBuffer);
        return ((int32_t)lRaw*1000000000L)+(lRaw>>32);
    }

    THeader ParseHeader()
    {
        uint32_t lLength=mReader.Get<uint32_t>();
        THeader lHeader;

        while(lLength)
        {
            uint32_t lFieldLength=mReader.Get<uint32_t>();
            if(lFieldLength>lLength)
                throw XData("invalid field length size");
            std::string lFieldName;
            get_line(mReader.Stream(),lFieldName,'=');
            auto lBuffer=read_simple_buffer(mReader.Stream(),lFieldLength-lFieldName.length()-1);
            lLength-=lFieldLength+4;
            lHeader.emplace(std::move(lFieldName),std::move(lBuffer));
        }

        return lHeader;
    }

    void ParseBagHeader(const THeader &pHeader)
    {
        //we don't care about index_pos
        mChunks.reserve(binary_interpret<int32_t>(pHeader.at("chunk_count")));
        mConnections.resize(binary_interpret<int32_t>(pHeader.at("conn_count")));
        mReader.Ignore(mReader.Get<uint32_t>()); //content is just filling
    }
    
    void ParseIndexData(const THeader &pHeader)
    {
        int32_t lConn=binary_interpret<int32_t>(pHeader.at("conn"));
        int32_t lCount=binary_interpret<int32_t>(pHeader.at("count"));

        if(lConn>=mConnections.size())
            throw XData("invalid conn in ros bag index data");
        
        uint32_t lLength=mReader.Get<uint32_t>();
        if(lLength!=12*lCount) throw XData("invalid index data field in ros bag");
        for(int i=0;i<lCount;i++)
        {
            SMessage lMessage;
            lMessage.mChunk=mChunks.size()-1;
            lMessage.mTimestamp=GetTimestamp(read_simple_buffer(mReader.Stream(),8));
            lMessage.mOffset=mReader.Get<int32_t>();
            mConnections[lConn].mMessages.push_back(lMessage);
        }
    }
    
    void ParseChunk(const THeader &pHeader)
    {
        std::string lCompression=binary_to_string(pHeader.at("compression"));
        SChunk lChunk;
        if(lCompression=="none")
        {
            lChunk.mCompressed=false;
        }
        else
        {
            ///\todo implement compression
            throw XNotImplemented("compressed chunks in ROS bags");
        }
        lChunk.mSize=mReader.Get<uint32_t>();
        lChunk.mOffset=mReader.Stream()->Tell();
        //just ignore the data for now, we're just reading headers
        mReader.Ignore(lChunk.mSize);
        mChunks.push_back(std::move(lChunk));
    }
    
    void ParseConnection(const THeader &pHeader)
    {
        int32_t lConn=binary_interpret<int32_t>(pHeader.at("conn"));
        if(lConn>=mConnections.size())
            throw XData("invalid connection index in ros bag");
        mConnections[lConn].mTopic=binary_to_string(pHeader.at("topic"));
        ///\todo maybe some of the info in data is useful
        THeader lExtra=ParseHeader();
        mConnections[lConn].mType=binary_to_string(lExtra.at("type"));
        mConnections[lConn].mDefinition=binary_to_string(lExtra.at("message_definition"));
    }
    
    bool ParseRecord()
    {
        THeader lHeader=ParseHeader();
        uint8_t lOp=binary_interpret<uint8_t>(lHeader.at("op"));
        
        switch(lOp)
        {
        case 3:
            ParseBagHeader(lHeader);
            break;
        case 4:
            ParseIndexData(lHeader);
            break;
        case 5:
            ParseChunk(lHeader);
            break;
        case 6:
            //chunk info, that's it!
            return false;
        case 7:
            ParseConnection(lHeader);
            break;
        default:
            msg_warn() << "unknown bag opcode " << lOp << "\n";
            mReader.Ignore(mReader.Get<uint32_t>());
        }
        return true;
    }
    
    const CSimpleBuffer<uint8_t> &GetChunkBuffer(SChunk &pChunk)
    {
        if(!pChunk.mBuffer)
        {
            mReader.Stream()->Seek(pChunk.mOffset);
            pChunk.mBuffer.reset(new CSimpleBuffer<uint8_t>(read_simple_buffer(mReader.Stream(),pChunk.mSize)));
        }
        
        return *pChunk.mBuffer;
    }
    
    CRange<const uint8_t*> GetMessageData(const SMessage &pMessage)
    {
        const auto &lBuffer=GetChunkBuffer(mChunks[pMessage.mChunk]);
        const uint8_t *lPtr=lBuffer.begin()+pMessage.mOffset;
        lPtr+=4+binary_interpret<uint32_t>(lPtr,lPtr+4);
        uint32_t lDataLength=binary_interpret<uint32_t>(lPtr,lPtr+4);
        lPtr+=4;
        return {lPtr,lPtr+lDataLength};
    }
    
    void ParseBag()
    {
        std::string lFileHeader;
        get_line(mReader.Stream(),lFileHeader);
        msg_info() << lFileHeader << "\n";
        start_timer();
        while(ParseRecord()) {}
        end_timer();
        
        msg_info() << "number of chunks:" << mChunks.size() << "\n";
        for(const auto &c : mConnections)
        {
            msg_info() << "\n";
            msg_info() << c.mTopic << "\n";
            msg_info() << c.mType << "\n";
            msg_info() << c.mDefinition << "\n";
            msg_info() << "   " << c.mMessages.size() << " messages\n";
//             if(c.mTopic!="/preprocessor/kinect1/depth") continue;
//             for(const auto &m : c.mMessages)
//             {
//                 io::CBinaryReader lReader(io::open_mem(GetMessageData(m)));
//                 lReader.Ignore(4); //seq
//                 lReader.Ignore(8); //timestamp
//                 lReader.Ignore(lReader.Get<uint32_t>()); //frame_id
//                 uint32_t lHeight=lReader.Get<uint32_t>();
//                 uint32_t lWidth=lReader.Get<uint32_t>();
//                 lReader.Ignore(lReader.Get<uint32_t>()); //encoding
//                 lReader.Ignore(1); //endianness
//                 uint32_t lStep=lReader.Get<uint32_t>();
//                 msg_info() << lWidth << "x" << lHeight << "/" << lStep << "\n";
//             }
            
            std::string lFolder;
            if(c.mTopic=="/preprocessor/kinect1/points")
                lFolder="topic1";
            if(c.mTopic=="/preprocessor/kinect2/points")
                lFolder="topic2";
            
            if(lFolder.empty()) continue;
            
            for(const auto &m : c.mMessages)
            {
                io::CBinaryReader lReader(io::open_mem(GetMessageData(m)));
                uint32_t lSeq=lReader.Get<uint32_t>();
                lReader.Ignore(8); //timestamp
                lReader.Ignore(lReader.Get<uint32_t>());
                uint32_t lHeight=lReader.Get<uint32_t>();
                uint32_t lWidth=lReader.Get<uint32_t>();
                uint32_t lNumFields=lReader.Get<uint32_t>();
                for(int i=0;i<lNumFields;i++)
                {
                    std::string lName=read_string(lReader.Stream(),lReader.Get<uint32_t>());
                    lReader.Ignore(9);
//                     uint32_t lOffset=lReader.Get<uint32_t>();
//                     uint8_t lType=lReader.Get<uint8_t>();
//                     uint32_t lCount=lReader.Get<uint32_t>();
//                     msg_info() << "field " << lName << " " << lOffset << " " << lType << " " << lCount << "\n";
                }
                lReader.Ignore(1);
                lReader.Ignore(12);
//                 uint32_t lPointStep=lReader.Get<uint32_t>();
//                 uint32_t lRowStep=lReader.Get<uint32_t>();
//                 uint32_t lSize=lReader.Get<uint32_t>();
                cvi::CCloud_rgb8 lCloud;
                for(int i=0;i<lWidth*lHeight;i++)
                {
                    cvi::CCloudPoint_rgb8 lPoint;
                    lPoint.mPoint.x()=lReader.Get<float>();
                    lPoint.mPoint.y()=lReader.Get<float>();
                    lPoint.mPoint.z()=lReader.Get<float>();
                    lReader.Ignore(4);
                    lPoint.mPixel[0]=lReader.Get<uint8_t>();
                    lPoint.mPixel[1]=lReader.Get<uint8_t>();
                    lPoint.mPixel[2]=lReader.Get<uint8_t>();
                    lReader.Ignore(13);
                    lCloud.push_back(lPoint);
                }
                cvi::save_cloud(lCloud,io::file(lFolder+"/frame"+to_string(lSeq)+".pcd"));
            }
        }
    }
    
private:
    io::CBinaryReader mReader;
    
    std::vector<SChunk> mChunks;
    std::vector<SConnection> mConnections;
};

int main(int pNArg,char **pArgs)
{
    buola_init(pNArg,pArgs);

    CROSBag lBag(io::arg(0));
    
    return buola_finish();
}
